#
# example005.py
#
# This example loads an audio wav file, DWT, next it plots their samples in time
# domain and its DFT (Discrete Fourier Transform) using FFT (Fast Fourier Trans-
# form)
#
# Copyright (C) 2012 Robert Buj Gelonch
# Copyright (C) 2012 David Megias Jimenez
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
__author__ = "Robert Buj Gelonch, and David Megias Jimenez"
__copyright__ = "Copyright 2012, Robert Buj Gelonch and David Megias Jimenez"
__credits__ = ["Robert Buj Gelonch", "David Megias Jimenez"]
__license__ = "GPL"
__version__ = "3"
__maintainer__ = "Robert Buj"
__email__ = "rbuj@uoc.edu"
__status__ = "Development"
__docformat__ = 'plaintext'

from numpy import arctan2
from numpy import imag
from numpy import log10
from numpy import ones
from numpy import real
from numpy import unwrap
from numpy import zeros
from pylab import figure
from pylab import plot
from pylab import show
from pylab import stem
from pylab import subplot
from pylab import subplots_adjust
from pylab import suptitle
from pylab import title
from pylab import xlabel
from pylab import ylabel
from pylab import ylim
from pylab import yscale
from pywt import dwt
from pywt import idwt
from scipy import arange
from scipy import fft
from scipy.io.wavfile import read
from scipy.signal import cascade
from scipy.signal import daub
from scipy.signal import freqz
from scipy.signal import lfilter

print "example005.py"
print

#--------------------------------------------------------------
# Reading data and obtaining basic attributes
#--------------------------------------------------------------
(sample_rate, samples) = read("../../../media/OSR_uk_000_0020_8k.wav")

#--------------------------------------------------------------
# Expand number of samples 2**|log2(n)|
#--------------------------------------------------------------
n = len(samples)
k = arange(n) # samples index vector
Ts = 1.0 / sample_rate
T = n / sample_rate
t = arange(0, n * Ts, Ts) # time vector

#--------------------------------------------------------------
# Frequency response
#--------------------------------------------------------------
frq = k / T # two sides frequency range
frq = frq[range(n / 2)] # one side frequency range
Y = fft(samples) / n # fft computing and normalization
Y = Y[range(n / 2)]

#--------------------------------------------------------------
# DWT - Daubechies wavelet
#--------------------------------------------------------------
numtaps = 10
daub_coeff = daub(numtaps)
(x, phi, psi) = cascade(daub_coeff)
cA, cD = dwt(samples, 'db10')
idwt_s = idwt(cA, cD, 'db10')

#--------------------------------------------------------------
# Low pass filter
#--------------------------------------------------------------
low_filtered_signal = lfilter(daub_coeff, 1.0, samples)
(w, h) = freqz(daub_coeff, 1.0)
h_dB = 20 * log10(abs(h))
h_Phase = unwrap(arctan2(imag(h), real(h)))

#--------------------------------------------------------------
# Frequency response LP filter
#--------------------------------------------------------------
# fft computing and normalization
YLP = fft(low_filtered_signal) / n
YLP = YLP[range(n / 2)]

#--------------------------------------------------------------
# Plotting
#--------------------------------------------------------------
fig = figure(num=None, \
             figsize=(16, 13), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
suptitle(r'Wavelet Transform Daubieches 10th order')


ax = subplot(421)
title('Scaling function')
ylim(-1.05, 1.05)
plot(phi)

ax = subplot(422)
title('Wavelet function')
ylim(-1.05, 1.05)
plot(psi)

ax = subplot(423)
title('cA')
plot(cA, 'b')
ax.set_xlim([0, len(cA)])

ax = subplot(424)
title('cD')
plot(cD)
ax.set_xlim([0, len(cD)])

ax = subplot(425)
title('cD/cA vs cA/cD, 1st 2,000 coeffs')
plot(abs(cA[0:2000] / cD[0:2000]), 'b')
plot(abs(cD[0:2000] / cA[0:2000]), 'r')
yscale('log')

ax = subplot(426)
title('IDWT(cA,cD) vs IDWT(cA,None)')
plot((idwt(cA, None, 'db10')), 'r')
plot((idwt_s), 'b')
ax.set_xlim([0, len(idwt_s)])

low_filtered_signal_down = zeros(len(low_filtered_signal) / 2)
for i in range(len(low_filtered_signal_down)):
    low_filtered_signal_down[i] = low_filtered_signal[2 * i]
ax = subplot(427)
title('cA vs down(LPF(s))')
plot(low_filtered_signal_down, 'r')
plot(cA, 'b')
ax.set_xlim([0, len(low_filtered_signal_down)])

ax = subplot(428)
title('Diff between Original Signal and IDWT(cA,cD)')
plot(abs(samples-idwt_s), 'b')
ax.set_xlim([0, len(idwt_s)])

# Save plot
fig.savefig("../plots/example005-py-1.png", format="png")

# Show plot
show()

#--------------------------------------------------------------
# Plotting low pass filter
#--------------------------------------------------------------
impulse = ones(50)
x = arange(0, 50)
response = lfilter(daub_coeff, 1.0, impulse)

fig = figure(num=None, \
             figsize=(14, 9), \
             dpi=80, \
             facecolor='w', \
             edgecolor='k')
suptitle(r'Low Pass Filter')
subplot(221)
stem(x, response)
title(r'Impulse response')

# Freq plot
subplot(223)
xlabel('Freq (Hz)')
ylabel('|Y(freq)|')
title('Spectogram')
plot(frq, abs(Y), 'b')
plot(frq, abs(YLP), 'r')

subplot(222)
plot(w / max(w), h_dB)
ylabel('Magnitude (db)')
xlabel(r'Normalized Frequency (x$\pi$rad/sample)')
title(r'Frequency response')

subplot(224)
plot(w / max(w), h_Phase)
ylabel('Phase (radians)')
xlabel(r'Normalized Frequency (x$\pi$rad/sample)')
title(r'Phase response')
subplots_adjust(hspace=0.5)

# Save plot
fig.savefig("../plots/example005-py-2.eps", format="eps")

# Show plot
show()

print "Done"